//
// Created by liugang on 9/12/25.
//

#ifndef VARS_H
#define VARS_H

#include "nodes.h"

#include <vector>
#include <tuple>
#include <array>
#include <cstdint>

// struct Dual {
//     double v;
//     double dv;
//     Dual(double v=0, double dv=0): v(v), dv(dv) {}
// };
//
// inline Dual operator+(Dual a, Dual b){ return {a.v+b.v, a.dv+b.dv}; }
// inline Dual operator-(Dual a, Dual b){ return {a.v-b.v, a.dv-b.dv}; }
// inline Dual operator*(Dual a, Dual b){ return {a.v*b.v, a.v*b.dv + a.dv*b.v}; }
// inline Dual operator/(Dual a, Dual b){
//     double inv = 1.0 / b.v;
//     return {a.v*inv, (a.dv*b.v - a.v*b.dv)*inv*inv};
// }
// inline Dual exp(Dual a){ double e=std::exp(a.v); return {e, e*a.dv}; }
//
// 作者：Reminder
// 链接：https://www.zhihu.com/question/1946816589988102594/answer/1949824781638703098
// 来源：知乎
// 著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。

// f = For(init, cond, itr)
//     f.stmt()
//     f.stmt()

class FLOW;

struct Var {
    FLOW* flow{};
    int32_t access_proc_id;
    int32_t loc_id{};

    explicit Var(FLOW* f, int32_t a, int32_t id)
        : flow(f), access_proc_id(a), loc_id(id) {
    }
};


/**
 * 语法盒子
 */
struct SyBox {
    Node* node{};

    // explicit SyBox() {
    //     node = new Stmt{NK_STMT_BLOCK, NF_BLOCK_SEQ}; // 空语句块，这样可以直接写 FOR({}, xxx,{})
    // }

    explicit SyBox(Node* n)
        : node(n) {
    }

    SyBox(real v) {
        node = new ExprLiteral(v);
    }

    SyBox(Var v) {
        node = new ExprVar(v.flow, v.access_proc_id, v.loc_id);
    }

    auto operator =(SyBox v) -> SyBox {
        Expr* assign = new Expr(NK_EXPR_ASSIGN);
        assign->add_child(node);
        assign->add_child(v.node);
        node = assign;
        return SyBox{assign};
    }
};

using AccessProc = real(*)(void* context, int32_t loc_id);

struct AccessVar {
    int32_t access_proc_id;
    int32_t loc_id{};
};


struct Access {
    /**
     * Access 函数的注册编号
     */
    int32_t access_proc_id;

    explicit Access(int32_t proc_id) {
        access_proc_id = proc_id;
    }

    /**
     * 提供了一个方法用来生成语法树
     */
    auto operator()(int32_t loc_id) -> AccessVar {
        return AccessVar{access_proc_id, loc_id};
    }

    // /**
    //  * 提供了一个方法用来生成语法树
    //  */
    // auto operator()(int32_t loc_id) -> SyBox {
    //     ExprAccess* access = new ExprAccess(flow, access_index, loc_id);
    //     return SyBox{access};
    // }
};

using EvalStmtHookProc = void(*)(void* ctx, const Node* curr);
using EvalExprHookProc = void(*)(void* ctx, const Node* curr, real value);

// 调试探针
struct DebugProbe {
    void* debug_context{nullptr};
    EvalStmtHookProc eval_stmt_hook{nullptr};
    EvalExprHookProc eval_expr_hook{nullptr};
};

class FLOW {
public:
    FLOW() {
        root_ = new Stmt(NK_STMT_BLOCK, NF_BLOCK_SEQ);
        curr_ = root_;
    }

    // auto BEGIN() -> FLOW& {
    //     auto block = new Stmt{NK_STMT_BLOCK};
    //     if (root_ == nullptr) [[unlikely]] {
    //         root_ = new Stmt{NK_STMT_BLOCK};
    //         curr_ = root_;
    //     } else {
    //         curr_->add_child(block);
    //     }
    // }
    //
    // auto END() -> FLOW& {
    //     curr_->add_child(new Stmt(NK_STMT_END));
    //     curr_ = curr_->parent;
    // }

    // FOR
    // +-- init_stmt
    //     +-- init_expr
    // +-- JUMP:END
    //     +-- cond_expr
    // +-- body_stmt
    // +-- iter_stmt
    //     +-- iter_expr
    // +-- JUMP[continue]
    // +-- END
    auto FOR(SyBox init_expr, SyBox cond_expr, SyBox iter_expr) -> FLOW& {
        // FOR(不参与求导)
        Stmt* for_stmt = new Stmt(NK_STMT_BLOCK, NF_BLOCK_FOR | NF_NONDIFFERENTIABLE);

        // 关键的节点
        Stmt* end_for = new Stmt(NK_STMT_END);

        // +-- init_stmt
        //     +-- init_expr
        for_stmt->add_child(new Stmt(NK_STMT_EXPR, 0, init_expr.node));

        // +-- JUMP:END
        //     +-- cond_expr
        StmtJump* cond_jump = new StmtJump(NF_JUMP_FOR_COND, end_for, cond_expr.node);
        for_stmt->add_child(cond_jump);

        // +-- body_stmt
        Stmt* body = new Stmt(NK_STMT_BLOCK, NF_BLOCK_SEQ);
        for_stmt->add_child(body);

        // +-- iter_stmt
        //     +-- iter_expr
        Stmt* iter_stmt = new Stmt(NK_STMT_EXPR, 0, iter_expr.node);
        for_stmt->add_child(iter_stmt);

        // +-- JUMP[continue] -> JUMP:END
        StmtJump* jmp_continue = new StmtJump(NF_JUMP_FOR_CONTINUE, cond_jump);
        for_stmt->add_child(jmp_continue);

        // +-- END
        for_stmt->add_child(end_for);

        curr_->add_child(for_stmt);
        curr_ = body; // 自动锚定到 for 的 body

        // break和continue会跨语句，这里通过break_stack_和continue_stack_来共享跳转位置信息
        break_stack_.emplace_back(end_for);
        continue_stack_.emplace_back(iter_stmt);
        return *this;
    }

    auto ENDFOR() -> FLOW& {
        // 回朔到 for 节点
        while ((curr_ != nullptr) && ((curr_->flags & NF_BLOCK_MASK) != NF_BLOCK_FOR)) {
            curr_ = curr_->parent;
        }

        // for 的body 部分要加 END
        ASSERT((curr_->flags & NF_BLOCK_MASK) == NF_BLOCK_FOR);
        curr_->child(2)->add_child(new Stmt(NK_STMT_END));

        // 跳到 for 的父级节点
        curr_ = curr_->parent;
        ASSERT(curr_ != nullptr);

        break_stack_.pop_back();
        continue_stack_.emplace_back();
        return *this;
    }

    // BREAK 不改变 curr_，它只影响运行时指令指针的位置
    auto BREAK() -> FLOW& {
        Node* break_target = break_stack_.back();
        ASSERT(break_target != nullptr);
        Stmt* break_stmt = new StmtJump(NF_JUMP_FOR_BREAK, break_target);
        curr_->add_child(break_stmt);
        return *this;
    }

    // BREAK 不改变 curr_，它只影响运行时指令指针的位置
    auto CONTINUE() -> FLOW& {
        Node* continue_target = continue_stack_.back();
        ASSERT(continue_target != nullptr);
        Stmt* continue_stmt = new StmtJump(NF_JUMP_FOR_CONTINUE, continue_target);
        curr_->add_child(continue_stmt);
        return *this;
    }

    // IF
    // +-- JUMP:END
    //     +-- cond_expr
    // +-- true_stmt
    // +-- END
    //
    // IF
    // +-- JUMP:false_stmt
    //     +-- cond_expr
    // +-- true_stmt
    // +-- JMP:END
    // +-- false_stmt
    // +-- END
    auto IF(SyBox cond_expr) -> FLOW& {
        // IF
        Stmt* if_stmt = new Stmt(NK_STMT_BLOCK, NF_BLOCK_IF);

        // 必须要有的end标记和false_stmt
        Stmt* if_end = new Stmt(NK_STMT_END);

        // +-- JUMP:END
        //     +-- cond_expr
        StmtJump* cond_jump = new StmtJump(NF_JUMP_IF_COND, if_end, cond_expr.node);
        if_stmt->add_child(cond_jump);

        // +-- true_stmt
        Stmt* true_stmt = new Stmt(NK_STMT_BLOCK, NF_BLOCK_SEQ);
        if_stmt->add_child(true_stmt);

        // +-- END
        if_stmt->add_child(if_end);

        curr_->add_child(if_stmt);
        curr_ = true_stmt;
        return *this;
    }

    auto ELSE() -> FLOW& {
        // 定位到最近的 if
        while ((curr_ != nullptr) && ((curr_->flags & NF_BLOCK_MASK) != NF_BLOCK_IF)) {
            curr_->add_child(new Stmt(NK_STMT_END));
            curr_ = curr_->parent;
        }

        ASSERT(curr_ != nullptr);
        ASSERT((curr_->flags & NF_BLOCK_MASK) == NF_BLOCK_IF);

        // 将 if 的结构找全
        // JUMP:END
        StmtJump* if_cond_jump = static_cast<StmtJump*>(curr_->child(0));
        Stmt* if_true_stmt = static_cast<Stmt*>(curr_->child(1));
        Stmt* if_end = static_cast<Stmt*>(curr_->child(2));

        // +-- JMP:END
        StmtJump* true_jmp_end = new StmtJump(NF_JUMP_FORCE, if_end);

        // 创建false分支语句块
        Stmt* false_stmt = new Stmt(NK_STMT_BLOCK, NF_BLOCK_SEQ);

        // 在 if_true_stmt 与 if_end 之间插入 true_jmp_end
        curr_->insert_child(if_true_stmt, if_end, static_cast<Stmt*>(true_jmp_end));

        // 在 true_jmp_end 与 if_end 之间插入 false_stmt
        curr_->insert_child(static_cast<Stmt*>(true_jmp_end), if_end, false_stmt);

        // 让 if_cond_jump 改为跳转到 false_stmt
        if_cond_jump->target = false_stmt;

        // 定位到最近的 if 的 false_stmt 分支
        curr_ = false_stmt;
        return *this;
    }

    auto ENDIF() -> FLOW& {
        // 定位到最近的 if
        while ((curr_ != nullptr) && ((curr_->flags & NF_BLOCK_MASK) != NF_BLOCK_IF)) {
            curr_->add_child(new Stmt(NK_STMT_END));
            curr_ = curr_->parent;
        }

        // 回到 if 的父级节点
        ASSERT(curr_ != nullptr);
        ASSERT((curr_->flags & NF_BLOCK_MASK) == NF_BLOCK_IF);
        curr_ = curr_->parent;
        return *this;
    }

    auto EXPR(SyBox expr) -> FLOW& {
        Stmt* expr_stmt = new Stmt(NK_STMT_EXPR, 0, expr.node);
        curr_->add_child(expr_stmt);
        return *this;
    }

    auto COMMIT() -> void {
        curr_->add_child(new Stmt(NK_STMT_STOP));
    }

    auto root() -> Node* {
        return root_;
    }

    auto root() const -> const Node* {
        return root_;
    }

    auto bind_target(void* inst, void* model) -> void {
        inst_ = inst;
        model_ = model;
    }

    auto bind_context(void* ctx) -> void {
        context_ = ctx;
    }

    auto var(std::string name, AccessVar accvar) -> Var {
        if (vars_.size() < (accvar.loc_id + 1)) {
            vars_.resize(accvar.loc_id + 1);
            var_names_.resize(accvar.loc_id + 1);
        }

        ASSERT((accvar.loc_id >= 0) && (accvar.loc_id < static_cast<int32_t>(vars_.size())));
        ASSERT((accvar.loc_id >= 0) && (accvar.loc_id < static_cast<int32_t>(var_names_.size())));
        ASSERT(vars_.size() == var_names_.size());

        var_names_[accvar.loc_id] = std::move(name);
        vars_[accvar.loc_id] = 0.0;
        ASSERT(var_names_.size() == vars_.size());
        return Var{this, accvar.access_proc_id, accvar.loc_id};
    }

    auto var_value(int32_t var_index) -> real& {
        ASSERT((var_index >= 0) && (var_index < static_cast<int32_t>(vars_.size())));
        return vars_[var_index];
    }

    auto var_value(int32_t var_index) const -> const real& {
        ASSERT((var_index >= 0) && (var_index < static_cast<int32_t>(vars_.size())));
        return vars_[var_index];
    }

    auto var_value(Var var) -> real& {
        return var_value(var.loc_id);
    }

    auto var_value(Var var) const -> const real& {
        return var_value(var.loc_id);
    }

    auto var_name(int32_t index) -> std::string {
        return var_names_[index];
    }

    auto reset_vars() -> void {
        for (auto& var : vars_) {
            var = 0.0;
        }
    }

    auto eval(void* context) -> void {
        ASSERT(nullptr);
        context_ = context;
    }

    auto access(std::string access_name, AccessProc proc) -> Access {
        int32_t access_index = static_cast<int32_t>(accesses_.size());
        access_names_.emplace_back(std::move(access_name));
        accesses_.emplace_back(proc);
        ASSERT(accesses_.size() == access_names_.size());
        return Access{access_index};
    }

    auto access(int32_t access_index) const -> AccessProc {
        ASSERT((access_index >= 0) && (access_index < accesses_.size()));
        return accesses_[access_index];
    }

    auto access_name(int32_t access_index) const -> std::string {
        ASSERT((access_index >= 0) && (access_index < access_names_.size()));
        return access_names_[access_index];
    }

    auto context() const -> void* {
        return context_;
    }

    /**
     * 安装调试探针
     */
    auto setup_debug_probe(DebugProbe* probe) -> void {
        ASSERT(probe != nullptr);
        probe_ = probe;
    }

    auto debug_probe() const -> DebugProbe* {
        return probe_;
    }

private:
    void* inst_{};
    void* model_{};
    void* context_{nullptr};
    std::vector<real> vars_;
    std::vector<AccessProc> accesses_;

    std::vector<std::string> var_names_;
    std::vector<std::string> access_names_;
    Node* root_{nullptr};
    Node* curr_{nullptr};
    std::vector<Node*> break_stack_;
    std::vector<Node*> continue_stack_;

    DebugProbe* probe_{nullptr};
};

static inline auto operator+(SyBox l, SyBox r) -> SyBox {
    Expr* node = new Expr(NK_EXPR_ADD);
    node->add_child(l.node);
    node->add_child(r.node);
    return SyBox{node};
}

static inline auto operator-(SyBox l, SyBox r) -> SyBox {
    Expr* node = new Expr(NK_EXPR_SUB);
    node->add_child(l.node);
    node->add_child(r.node);
    return SyBox{node};
}

static inline auto operator*(SyBox l, SyBox r) -> SyBox {
    Expr* node = new Expr(NK_EXPR_MULT);
    node->add_child(l.node);
    node->add_child(r.node);
    return SyBox{node};
}

static inline auto operator/(SyBox l, SyBox r) -> SyBox {
    Expr* node = new Expr(NK_EXPR_DIV);
    node->add_child(l.node);
    node->add_child(r.node);
    return SyBox{node};
}

static inline auto operator>(SyBox l, SyBox r) -> SyBox {
    Expr* node = new Expr(NK_EXPR_GT, NF_NONDIFFERENTIABLE);
    node->add_child(l.node);
    node->add_child(r.node);
    return SyBox{node};
}

static inline auto operator<(SyBox l, SyBox r) -> SyBox {
    Expr* node = new Expr(NK_EXPR_LT, NF_NONDIFFERENTIABLE);
    node->add_child(l.node);
    node->add_child(r.node);
    return SyBox{node};
}

static inline auto operator==(SyBox l, SyBox r) -> SyBox {
    Expr* node = new Expr(NK_EXPR_EQ, NF_NONDIFFERENTIABLE);
    node->add_child(l.node);
    node->add_child(r.node);
    return SyBox{node};
}

static inline auto sin(SyBox l) -> SyBox {
    Expr* node = new Expr(NK_EXPR_SIN);
    node->add_child(l.node);
    return SyBox{node};
}

static inline auto cos(SyBox l) -> SyBox {
    Expr* node = new Expr(NK_EXPR_COS);
    node->add_child(l.node);
    return SyBox{node};
}

static inline auto ref(real& r) -> SyBox {
    ExprRef* node = new ExprRef(r);
    return SyBox{node};
}

// static inline auto val(real& r) -> SyBox {
//     ExprVal* node = new ExprVal(r);
//     return SyBox{node};
// }

#endif //VARS_H